from importlib.resources import path
import os, sys
import shutil
from tqdm import tqdm
import random
import json
import cv2
import pandas as pd
from pathlib import Path
import glob
sys.path.append('..')
from utils.json import json_to_image_one_by_one
from utils.helper import read_xml, find_nodes, change_node_text, indent, write_xml


img_root_path = r'/root/project/AutoRepair_T7/data/t7_testing_data/t7_testing_data/all_cut_pad'
mask_root_path = r'/root/project/AutoRepair_T7/data/t7_testing_data/t7_testing_data/all_cut_pad_seg_predict/pred_mask'
dst_path = r'/root/project/AutoRepair_T7/data/t7_testing_data/t7_testing_data/all_img_cut_pad_using_pred_mask'
h, w = 224, 224  # crop size


# 正常code为5个字母/数字，后面可能带一个数字(第六位)
# 数字为1表示前程切割
# 数字为2表示长线，当前方案是去掉该数字
# 数字3表示缺陷衍生残留，当前方案是去掉该数字
# def label_change(label):
#     if len(label) == 6 and label[-1] in ['2', '3']:
#         label = label[:-1]
#     return label


def mask_2_box(mask_path):
    code = str(Path(mask_path).parent.name)
    img_path = os.path.join(img_root_path, code, str(Path(mask_path).stem) + '.jpg')
    img_ori = cv2.imread(img_path, 1)
    mask = cv2.imread(mask_path, 0)

    mask[mask > 0] = 128
    mask = mask.astype('uint8')
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    # for index, (label, mask) in enumerate(json_to_image_one_by_one(img_shape=img_ori.shape, 
    #                          json_path=json_path, toImage=False, visual=False, mask_label='box')):
    #     label = label_change(label)
    #     mask[mask > 0] = 128
    #     mask = mask.astype('uint8')
    #     contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    for index,cnt in enumerate(contours):
        (x,y), radius = cv2.minEnclosingCircle(cnt)

        if int(y-int(h/2)+1) < 0:
            y = y + (0 - int(y-int(h/2)+1))
        if int(x-int(w/2)+1) < 0:
            x = x + (0 - int(x-int(w/2)+1))
        if int(y+int(h/2)+1) > img_ori.shape[0]:
            y = y - (int(y+int(h/2)+1) - img_ori.shape[0])
        if int(x+int(w/2)+1) > img_ori.shape[1]:
            x = x - (int(x+int(w/2)+1) - img_ori.shape[1])
        y_min = int(y-int(h/2)+1)
        x_min = int(x-int(w/2)+1)
        y_max = int(y+int(h/2)+1)
        x_max = int(x+int(w/2)+1)

        # img_boxed = img_ori[y_min:y_max, x_min:x_max, ...]

        # boxed_out = os.path.join(dst_path, 'images', code)
        # os.makedirs(boxed_out, exist_ok=True)
        # cv2.imwrite(os.path.join(boxed_out, Path(img_path).stem + '_%d.png'%(index+1)), img_boxed)

        base_tree = read_xml("./base_example.xml")
        root = base_tree.getroot()
        anno_tree = read_xml("./anno_example.xml")
        folder_node = find_nodes(base_tree, "folder")
        filename_node = find_nodes(base_tree, "filename")
        path_node = find_nodes(base_tree, "path")
        width_node = find_nodes(base_tree, "size/width")
        height_node = find_nodes(base_tree, "size/height")
        depth_node = find_nodes(base_tree, "size/depth")
        change_node_text(folder_node, code)
        change_node_text(filename_node, str(Path(img_path).name))
        change_node_text(path_node, img_path)
        change_node_text(width_node, str(img_ori.shape[0]))
        change_node_text(height_node, str(img_ori.shape[1]))
        change_node_text(depth_node, str(img_ori.shape[2]))

        xmin_node = find_nodes(anno_tree, "bndbox/xmin")
        ymin_node = find_nodes(anno_tree, "bndbox/ymin")
        xmax_node = find_nodes(anno_tree, "bndbox/xmax")
        ymax_node = find_nodes(anno_tree, "bndbox/ymax")
        change_node_text(xmin_node, str(x_min))
        change_node_text(ymin_node, str(y_min))
        change_node_text(xmax_node, str(x_max))
        change_node_text(ymax_node, str(y_max))
        root.append(anno_tree.getroot())
        indent(root)

        label_out = os.path.join(dst_path, 'annotations', code)
        os.makedirs(label_out, exist_ok=True)
        write_xml(base_tree, os.path.join(label_out, Path(img_path).stem + '_%d.xml'%(index+1)))

    #     if len(radius_list) > 0:
    #         index = np.argmax(radius_list)
    #         # print(x_list[index], y_list[index])
    #         w = args.crop_size
    #         h = args.crop_size
    #         x = int(x_list[index])
    #         y = int(y_list[index])
    #         if int(y-int(h/2)+1) < 0:
    #             y = y + (0 - int(y-int(h/2)+1))
    #         if int(x-int(w/2)+1) < 0:
    #             x = x + (0 - int(x-int(w/2)+1))
    #         if int(y+int(h/2)+1) > img_ori.shape[0]:
    #             y = y - (int(y+int(h/2)+1) - img_ori.shape[0])
    #         if int(x+int(w/2)+1) > img_ori.shape[1]:
    #             x = x - (int(x+int(w/2)+1) - img_ori.shape[1])
    #         y_min = int(y-int(h/2)+1)
    #         x_min = int(x-int(w/2)+1)
    #         y_max = int(y+int(h/2)+1)
    #         x_max = int(x+int(w/2)+1)

    #         img_boxed = img_ori[y_min:y_max, x_min:x_max, ...]
    #         assert img_boxed.shape==(args.crop_size,args.crop_size,3), f'The size is {img_boxed.shape}, x is {x} and y is {y}.'
    #         cv2.imwrite(os.path.join(boxed_out, names[0] + '.png'), img_boxed)

    #         xmin_node = find_nodes(anno_tree, "bndbox/xmin")
    #         ymin_node = find_nodes(anno_tree, "bndbox/ymin")
    #         xmax_node = find_nodes(anno_tree, "bndbox/xmax")
    #         ymax_node = find_nodes(anno_tree, "bndbox/ymax")
    #         change_node_text(xmin_node, str(x_min))
    #         change_node_text(ymin_node, str(y_min))
    #         change_node_text(xmax_node, str(x_max))
    #         change_node_text(ymax_node, str(y_max))
    #         root.append(anno_tree.getroot())
    #         indent(root)
    #         write_xml(base_tree, os.path.join(label_out, names[0] + '.xml'))

    # with open(json_path, 'r') as f:
    #     content = f.read()
    #     data = json.loads(content)
    
    # data['imageData'] = 'null'
    
    # new_data_shapes= []
    # for d in data['shapes']:
    #     if len(d['points']) < 3:
    #         continue
    #     for i,p in enumerate(d['points']):  #[w, h]
    #         d['points'][i][0] = min(p[0], w)
    #         d['points'][i][1] = min(p[1], h)
    #     new_data_shapes.append(d)
    # data['shapes'] = new_data_shapes

    # with open(json_path, 'w') as f:
    #     b = json.dumps(data, indent=4)
    #     f.write(b)

def main():

    df = pd.DataFrame(data={'mask':glob.glob(os.path.join(mask_root_path, '*/*.png'))})  # 获取路径下所有json文件
    df = df.sample(frac=1, random_state=10).reset_index(drop=True)  # 打乱DataFrame

    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False
    
    if is_pandarallel:
        df['mask'].parallel_apply(lambda x: mask_2_box(x))
    else:
        df['mask'].apply(lambda x: mask_2_box(x))
    print()


if __name__=='__main__':
    main()
